import datasets
from pprint import pprint
import os
import sys
sys.path.insert(0, os.path.dirname(__file__) + "/..")
from utils.base import setup_env, setup_workdir

# 打印huggingface数据集情况
def list_datasets() -> None:
    my_datasets_list = datasets.list_datasets()
    print("Number of datasets in the Datasets library: ", len(my_datasets_list))
    pprint(my_datasets_list[:20], compact=True)  # 打印数据集列表中的前 20 个来看看


# 打印指定数据集的详细信息
def print_dataset_info(dataset_name: str) -> None:
	my_datasets_list = datasets.list_datasets()
	index = my_datasets_list.index(dataset_name)  # 获取该数据集在 my_datasets(其 type 为list) 的位置
	# 通过设置 with_details 为 True 来获得数据集的详细信息
	info = datasets.list_datasets(with_details=True)[index]
	pprint(info, compact=True)

# 下载指定数据集，保存到本地
def download_dataset(dataset_name: str, cache_dir: str ='./cache', split: str = None, train_test_split: float = 0.3) -> datasets.DatasetDict:
	# 通过 load_dataset() 函数来加载数据集，该函数的参数为数据集的名称，该名称可以是数据集的名称，也可以是数据集的 ID
	my_dataset_all = datasets.load_dataset(path=dataset_name, cache_dir=cache_dir, 
                                        split=split, download_mode = 'force_redownload',
                                        ignore_verifications=False)
	# 通过 train_test_split() 函数来划分数据集，该函数的参数为划分比例，该函数会返回划分后的数据集，其中包含了 train和test 2个部分
	if split != None:
		my_dataset_all = my_dataset_all.train_test_split(test_size=train_test_split)

	# 打印数据集的信息，逐个打印 train、validation 和 test 三个部分的信息
	for(split_name, split_dataset) in my_dataset_all.items():
		print("split_name:", split_name)
		print("split_dataset:", split_dataset)
	# 通过 save_to_disk() 函数来保存数据集到磁盘，保存的路径为 ./save_data，该路径下会生成一个名为 ChnSentiCorp 的文件夹，里面包含了数据集的所有信息
	my_dataset_all.save_to_disk(dataset_dict_path=dataset_name + '/save_data')
	return my_dataset_all

# 从磁盘加载数据集, 数据采用arrow格式存储
def load_dataset_from_disk(dataset_name: str) -> datasets.DatasetDict:
	# 通过 load_from_disk() 函数来加载数据集，该函数的参数为数据集的路径
	my_dataset = datasets.load_from_disk(dataset_path=dataset_name + '/save_data')
	print("my_dataset:", my_dataset)
	return my_dataset

# 加载本地数据集
def load_local_dataset(dataset_name: str, data_files: str, pattern: str = 'csv', cache_dir: str ='./cache', split: str = None, train_test_split: float = 0.3) -> datasets.DatasetDict:
	# 通过 load_dataset() 函数来加载数据集，该函数的参数为数据集的路径
	my_dataset_all = datasets.load_dataset(path=pattern, data_files=data_files, cache_dir=cache_dir, 
										split=split, ignore_verifications=False)
	# 通过 train_test_split() 函数来划分数据集，该函数的参数为划分比例，该函数会返回划分后的数据集，其中包含了 train和test 2个部分
	if train_test_split > 0:
		my_dataset_all = my_dataset_all.train_test_split(test_size=train_test_split)
	
	# 打印数据集的信息，逐个打印 train、validation 和 test 三个部分的信息
	for(split_name, split_dataset) in my_dataset_all.items():
		print("split_name:", split_name)
		print("split_dataset:", split_dataset)
	# 通过 save_to_disk() 函数来保存数据集到磁盘，保存的路径为 ./save_data，该路径下会生成一个名为 ChnSentiCorp 的文件夹，里面包含了数据集的所有信息
	my_dataset_all.save_to_disk(dataset_dict_path=dataset_name + '/save_data')
	return my_dataset_all


#%%
def test_remote_dataset():
	setup_env()
	setup_workdir(os.path.dirname(__file__))
	dataset_name = 'seamew/ChnSentiCorp'
	# list_datasets()
	# 判断数据集是否存在, 目录save_data下是否有dataset_dict.json文件
	if os.path.exists(dataset_name) == True and os.path.exists(dataset_name + '/save_data/dataset_dict.json') == True:
		my_dataset = load_dataset_from_disk(dataset_name)
	else:
		# print_dataset_info(dataset_name)
		# my_dataset = download_dataset(dataset_name, cache_dir='./cache')
		my_dataset = download_dataset(dataset_name, cache_dir='./cache', split='train[:100]')
	return my_dataset

#%%
def test_local_dataset()-> datasets.DatasetDict:
	setup_env()
	setup_workdir(os.path.dirname(__file__))
	config = {'dataset_name':'imdb', 'data_files': '../../IMDB-sentiment-analysis/IMDB-Dataset-zhch.csv'}
	# 判断数据集是否存在, 目录save_data下是否有dataset_dict.json文件
	dataset_name = config['dataset_name']
	if os.path.exists(dataset_name) == True and os.path.exists(dataset_name + '/save_data/dataset_dict.json') == True:
		my_dataset = load_dataset_from_disk(dataset_name)
	else:
		my_dataset = load_local_dataset(dataset_name, config['data_files'], cache_dir='./cache', split='train[:100]')
	return my_dataset
    
#%%
if __name__ == "__main__":
	# test_remote_dataset()
	test_local_dataset()	